from __future__ import print_function

import os
import numpy as np
import sklearn.utils


def load_train_data(shuffle=False):
    curr_dir = os.path.dirname(os.path.abspath(__file__))
    curr_dir = curr_dir.replace('\\', '/')

    X = np.load(curr_dir + '/npy_files/X_train.npy')
    y = np.load(curr_dir + '/npy_files/y_train.npy')
    m = np.load(curr_dir + '/npy_files/m_train.npy')

    if shuffle:
        X, y, m = sklearn.utils.shuffle(X, y, m, random_state=0)

    return X, y, m


def load_validation_data():
    curr_dir = os.path.dirname(os.path.abspath(__file__))
    curr_dir = curr_dir.replace('\\', '/')

    X = np.load(curr_dir + '/npy_files/X_validate.npy')
    ids = np.load(curr_dir + '/npy_files/ids_validate.npy')
    m = np.load(curr_dir + '/npy_files/m_validate.npy')

    return ids, X, m


def split_data(X, y, m, split_ratio=0.1):
    """
    Split data into training and testing
    @split_ratio: Ratio of testing, e.g. 0.1 will produce 90% training, 10% testing
    """
    split = X.shape[0] * split_ratio

    X_train = X[split:, :, :, :]
    X_test = X[:split, :, :, :]

    y_train = y[split:, :]
    y_test = y[:split, :]

    m_train = m[split:, :]
    m_test = m[:split, :]

    return X_train, y_train, X_test, y_test, m_train, m_test

